import contextlib
import logging
from collections.abc import AsyncIterator
import click
import httpx
import mcp.types as types
from mcp.server.lowlevel import Server
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from starlette.applications import Starlette
from starlette.routing import Mount
from starlette.types import Scope, Receive, Send

OPENWEATHER_URL = "https://api.openweathermap.org/data/2.5/weather"
DEFAULT_UNITS = "metric"  # 摄氏度
DEFAULT_LANG = "zh_cn"


async def fetch_weather(city: str, api_key: str) -> dict[str, str]:
    params = {
        "q": city,
        "appid": api_key,
        "units": DEFAULT_UNITS,
        "lang": DEFAULT_LANG,
    }

    async with httpx.AsyncClient(timeout=10) as client:
        r = await client.get(OPENWEATHER_URL, params=params)
        r.raise_for_status()
        data = r.json()

    weather_main = data["weather"][0]["main"]
    description = data["weather"][0]["description"]
    temp = data["main"]["temp"]
    feels_like = data["main"]["feels_like"]
    humidity = data["main"]["humidity"]

    return {
        "city": city,
        "weather": weather_main,
        "description": description,
        "temp": f"{temp}°C",
        "feels_like": f"{feels_like}°C",
        "humidity": f"{humidity}%"
    }


@click.command()
@click.option("--port", default=3000, help="Listen port")
@click.option("--api-key", required=True, envvar="OPENWEATHER_API_KEY", help="OpenWeather API")
@click.option("--log-level", default="INFO", help="Log level(DEBUG|INFO|WARNING|ERROR|CRITICAL)")
@click.option("--json-response", is_flag=True, default=False, help="启用 JSON 响应而非 SSE 流")
def main(port: int, api_key: str, log_level: str, json_response: bool) -> int:
    logging.basicConfig(
        level=getattr(logging, log_level.upper()),
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    )
    logger = logging.getLogger("weather-server")

    app = Server("mcp-streamable-http-weather")

    @app.call_tool()
    async def call_tool(name: str, arguments: dict) -> list[types.TextContent]:
        ctx = app.request_context
        city = arguments.get("location")
        if not city:
            raise ValueError("参数中必须包含 'location'")

        await ctx.session.send_log_message(
            level="info",
            data=f"正在获取 {city} 的天气...",
            logger="weather",
            related_request_id=ctx.request_id
        )

        try:
            weather = await fetch_weather(city, api_key)
        except Exception as err:
            await ctx.session.send_log_message(
                level="error",
                data=str(err),
                logger="weather",
                related_request_id=ctx.request_id
            )
            raise

        await ctx.session.send_log_message(
            level="info",
            data="天气数据获取成功!",
            logger="weather",
            related_request_id=ctx.request_id
        )

        summary = f"{weather['city']}: {weather['description']}，温度 {weather['temp']}，体感 {weather['feels_like']}，湿度 {weather['humidity']}"
        return [types.TextContent(type="text", text=summary)]

    @app.list_tools()
    async def list_tools() -> list[types.Tool]:
        return [types.Tool(
            name="get-weather",
            description="查询指定城市的实时天气（OpenWeather）",
            inputSchema={
                "type": "object",
                "required": ["location"],
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "城市英文名，如 'Beijing'"
                    }
                }
            }
        )]

    session_manager = StreamableHTTPSessionManager(
        app=app,
        event_store=None,
        json_response=json_response,
        stateless=True
    )

    async def handle_streamable_http(scope: Scope, receive: Receive, send: Send) -> None:
        await session_manager.handle_request(scope, receive, send)

    @contextlib.asynccontextmanager
    async def lifespan(app: Starlette) -> AsyncIterator[None]:
        async with session_manager.run():
            logger.info("天气 MCP 服务器已启动！")

            try:
                yield
            finally:
                logger.info("天气服务器正在关闭...")

    starlette_app = Starlette(
        debug=False,
        routes=[Mount("/mcp", handle_streamable_http)],
        lifespan=lifespan
    )

    import uvicorn

    uvicorn.run(starlette_app, host="0.0.0.0", port=port)

    return 0


if __name__ == "__main__":
    main()
